# The data set used in this example is from http://archive.ics.uci.edu/ml/datasets/Wine+Quality
# P. Cortez, A. Cerdeira, F. Almeida, T. Matos and J. Reis.
# Modeling wine preferences by data mining from physicochemical properties. In Decision Support Systems, Elsevier, 47(4):547-553, 2009.

import os
import warnings
import sys
import configparser
from optparse import OptionParser

import pandas as pd
import numpy as np
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from sklearn.model_selection import train_test_split
from sklearn.linear_model import ElasticNet
from urllib.parse import urlparse
import mlflow
import mlflow.sklearn


import logging

logging.basicConfig(level=logging.WARN)
logger = logging.getLogger(__name__)


def eval_metrics(actual, pred):
    '''
    Calculate rmse, mae, r2
    '''
    rmse = np.sqrt(mean_squared_error(actual, pred))
    mae = mean_absolute_error(actual, pred)
    r2 = r2_score(actual, pred)
    return rmse, mae, r2

def main(args):

    config = configparser.ConfigParser()                                    
    config.read(cfg_file)
    config.sections()

    # Set properties of parameters inside config file.
    mode = 'normal'                 # initialization
    for key in config['main']:
        if key == 'alpha' :
            alpha_ = config.getfloat('main', key)    
        if key == 'l1ratio' :
            l1Ratio_ = config.getfloat('main', key)   
        if key == 'mode': 
            mode_ = config.get('main', key)
        if key == 'experimentid':
            experiment_ = config.get('main', key)

    alpha = float(sys.argv[1]) if len(sys.argv) > 1 else alpha_
    l1Ratio = float(sys.argv[2]) if len(sys.argv) > 2 else l1Ratio_
    mode = sys.argv[3] if len(sys.argv) > 3 else mode_
    experimentname = sys.argv[4] if len(sys.argv) > 4 else experiment_
    # For normal mode, system will test alpha dna l2_ratio value provided.
    # For loop mode, we will go through possible value of alpha and l1_ratio from 0 to 1
    if mode == 'normal':
        alphaList = [alpha]
        l1RatioList = [l1Ratio]
    elif mode == 'loop':
        alphaList = [0.2*x for x in range(5)]
        l1RatioList = [0.2*x for x in range(5)]
    
    print("Mode: {0}".format(mode))
    warnings.filterwarnings("ignore")       # Some warnings will be ignored.
    # np.random.seed(40)

    # Create a new experiment.
    experiment_id = mlflow.create_experiment(experimentname)
    experiment = mlflow.get_experiment(experiment_id)

    # Read the wine-quality csv file from the URL
    csv_url = ("http://archive.ics.uci.edu/ml/machine-learning-databases/wine-quality/winequality-red.csv")
    try:
        data = pd.read_csv(csv_url, sep=";")
    except Exception as e:
        logger.exception("Unable to download training & test CSV, check your internet connection. Error: %s", e)

    # Split the data into training and test sets. (0.75, 0.25) split.
    train, test = train_test_split(data)
    # The predicted column is "quality" which is a scalar from [3, 9]
    train_x = train.drop(["quality"], axis=1)
    test_x = test.drop(["quality"], axis=1)
    train_y = train[["quality"]]
    test_y = test[["quality"]]

    step = 0
    r2Final = 100
    alphaFinal = alphaList[0]
    l1RatioFinal = l1RatioList[0]
    print("Test Elasticnet model with different alpha and l1_ratio values.")
    #with mlflow.start_run(experiment_id=experiment_id):
    for alpha in alphaList:
        for l1_ratio in l1RatioList:
            # Generate a new run in the beginning of each loop, and end this run at the end of each loop.
            with mlflow.start_run(run_name="ElasticnetModelTest{0}".format(step), experiment_id=experiment_id):
                lr = ElasticNet(alpha=alpha, l1_ratio=l1_ratio, random_state=42)
                lr.fit(train_x, train_y)
                predicted_qualities = lr.predict(test_x)
                (rmse, mae, r2) = eval_metrics(test_y, predicted_qualities)
                rmse = round(rmse, 2)
                mae = round(mae, 2)
                r2 = round(r2, 2)
                print("Test - alpha: {0}, l1_ratio: {1}. Result: RMSE: {2}, MAE: {3}, R2: {4}".format(\
                    alpha, l1_ratio, rmse, mae, r2))
                metrics = {"alpha": alpha, "l1_ratio":l1_ratio, "rmse": rmse, "mae": mae, "r2": r2}
                mlflow.log_metrics(metrics)
                # Model registry does not work with file store
                tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme
                # Save this trained model under current run
                if tracking_url_type_store != "file":
                    mlflow.sklearn.log_model(lr, "model", registered_model_name="ElasticnetWineModel")
                else:
                    mlflow.sklearn.log_model(lr, "model")
                if r2 <= r2Final:
                    r2Final = r2
                    alphaFinal = alpha
                    l1RatioFinal = l1_ratio
            step += 1
            mlflow.end_run()    # end this run at the end of each loop.

    print("Best Elasticnet model (alpha=%f, l1_ratio=%f):" % (alphaFinal, l1RatioFinal))
    lr = ElasticNet(alpha=alphaFinal, l1_ratio=l1RatioFinal, random_state=42)
    lr.fit(train_x, train_y)
    predicted_qualities = lr.predict(test_x)
    (rmse, mae, r2) = eval_metrics(test_y, predicted_qualities)
    print(" Best model - alpha: {0}, l1_ratio: {1}. Result: RMSE: {2}, MAE: {3}, R2: {4}".format(\
            alpha, l1_ratio, rmse, mae, r2))
    mlflow.log_param("alpha", alpha)
    mlflow.log_param("l1_ratio", l1_ratio)

    tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme

    # Model registry does not work with file store
    if tracking_url_type_store != "file":
        # Register the model
        # There are other ways to use the Model Registry, which depends on the use case,
        # please refer to the doc for more information:
        # https://mlflow.org/docs/latest/model-registry.html#api-workflow
        mlflow.sklearn.log_model(lr, "model", registered_model_name="ElasticnetWineModel")
    else:
        mlflow.sklearn.log_model(lr, "model")

def parse_args():
    '''
    Parse and validate input arguments
    '''
    global cfg_file

    parser = OptionParser()
    current_dir = os.path.dirname(os.path.realpath(__file__))
    configfile_dir = os.path.abspath(os.path.join(current_dir,"./config.txt"))
    parser.add_option("-c", "--cfg-file", dest="cfg_file", default=configfile_dir,
                  help="Set the adaptor config file.", metavar="FILE")

    (options, _) = parser.parse_args()
    cfg_file = options.cfg_file
    return 0

if __name__ == '__main__':
    ret = parse_args()
    #If argumer parsing fail, return failure (non-zero)
    if ret == 1:
        sys.exit(1)
    sys.exit(main(sys.argv))

